library(dplyr)
library(tidyr)
library(stringr)
library(ggplot2)
library(ggpubr)
library(ggrepel)
library(viridis)

if (!file.exists("output")) {
  dir.create("output")
}

# load data

source("HNC_metadata_tidy.R")
source("data_handling.R")

sbsCOSMIC <- read.csv("Input_signature_attributions/Input_SBS_COSMIC_attributions.csv", stringsAsFactors = T, row.names = 1) %>% tibble::rownames_to_column("donor_id")
dbsCOSMIC <- read.csv("Input_signature_attributions/Input_DBS_COSMIC_attributions.csv", stringsAsFactors = T, row.names = 1) %>% tibble::rownames_to_column("donor_id")
idCOSMIC <- read.csv("Input_signature_attributions/Input_ID_COSMIC_attributions.csv", stringsAsFactors = T, row.names = 1) %>% tibble::rownames_to_column("donor_id")

additional_df <- list(sbsCOSMIC, dbsCOSMIC, idCOSMIC)
for (df in additional_df) {
  data <- data %>% left_join(df, by = "donor_id")
}

data <- data %>% mutate(
  SBS_tobacco = SBS4 + SBS92 + SBS_I,
  DBS_tobacco = DBS2 + DBS6,
  ID_tobacco = ID3
)

# Dichotomize the signatures (Presence (signatures>0)="1",absence (signatures==0)="0")
# If signature is present in more than 75% cases, dichotomize by the median
data <- dichotomizeSigs(metadata = data, sigsn = data[c("donor_id", "SBS_tobacco", "DBS_tobacco", "ID_tobacco", "SBS4", "SBS92", "SBS_I", "DBS2", "ID3", "DBS6")])

countries <- c("Argentina", "Brazil", "Colombia", "Czechia", "Greece", "Italy", "Romania", "Slovakia")
files <- list.files("GLOBOCAN_2022/", pattern = "ASR")

asr_sub <- NULL
for (f in files) {
  asr_f <- read.csv(paste0("GLOBOCAN_2022/", f), row.names = NULL) %>%
    select("Label", "ASR..World.") %>%
    rename(Population = `Label`, ASR = `ASR..World.`) %>%
    filter(Population %in% countries) %>%
    mutate(
      Population = ifelse(Population == "Czechia", "Czech Republic", Population),
      subsite = str_split_1(f, "[_.]")[2],
      sex = str_split_1(f, "[_.]")[3],
      code = paste(Population, subsite, paste0(sex, "s"))
    )
  asr_sub <- rbind(asr_sub, asr_f)
}

data <- data %>%
  mutate(
    code = paste(country, subsite, paste0(sex, "s")),
    age = scale(age_diag)
  ) %>%
  left_join(asr_sub[c("code", "ASR")], by = "code")

confounders <- "age_group"

signatures <- c("SBS_tobacco", "DBS_tobacco", "ID_tobacco")

output <- NULL
for (s in signatures) {
  sigsn <- ifelse(data[, s] > 0, 1, 0)
  if (sum(sigsn / nrow(data)) > .75) {
    # linear regression
    myformula <- paste0(s, " ~ ", paste(c("ASR", confounders), collapse = " + "))
    model <- lm(myformula, data = data)
    result <- summary(model)$coefficients
  } else {
    # logistic regression
    myformula <- paste0(paste0(s, "_cat"), " ~ ", paste(c("ASR", confounders), collapse = " + "))
    model <- glm(myformula, data = data, family = binomial("logit"))
    result <- summary(model)$coefficients
  }

  # build output
  r <- data.frame(
    signature = s,
    pval = result[rownames(result) == "ASR"][4],
    p_adj = ifelse(result[rownames(result) == "ASR"][4] * length(signatures) < 1, result[rownames(result) == "ASR"][4] * length(signatures), 1),
    Estimate = result[rownames(result) == "ASR"][1]
  )
  output <- rbind(output, r)

  country_burden <- data %>%
    group_by(code) %>%
    summarize(
      n = n(),
      ASR = median(ASR),
      burden = mean(get(s)),
      freq = sum(as.integer(as.character(get(paste0(s, "_cat")))))
    ) %>%
    mutate(freq = freq / n * 100)

  # plot
  p <- ggplot(data, aes(x = ASR, y = get(s))) +
    geom_smooth(method = lm, colour = "darkblue", fill = "lightblue", size = 0.2) +
    geom_point(data = country_burden, aes(x = ASR, y = burden, size = n, color = freq), alpha = .7) +
    geom_text_repel(data = country_burden, aes(x = ASR, y = burden), label = country_burden$code, size = 3.5, max.overlaps = 4) +
    labs(
      title = paste0("ASR/", str_replace_all(s, "_", " "), " associations across countries"),
      subtitle = paste("p =", signif(r$pval, 4)),
      x = "Incidence (ASR)", y = paste("Mutation burden"),
      color = "Frequency", size = "Number of samples"
    ) +
    scale_color_viridis() +
    theme_minimal() +
    theme(
      legend.position = "none",
      plot.title = element_text(face = "bold", size = 12),
      plot.subtitle = element_text(size = 10),
      axis.text = element_text(size = 12),
      axis.line = element_line(size = .3,color = "#404040"),
      axis.ticks = element_line(size = .3,color = "#404040"),
      panel.grid = element_blank()
    )

  print(p)
  ggsave(
    plot = p,
    filename = paste0("output/Figure_4c_", s, "_", Sys.Date(), ".pdf"),
    device = "pdf", width = 4.4, height = 4.5, units = "in"
  )
}

output
